import math
import numpy as np
from model.search.LocalFeatureSet import LocalFeatureSet
from os.path import join
import os
from model.common.common import *
import logging

    # selected_keypoints = keypoints[result_list,:]
    # # for i in range(len(result_list)):
        
    # #     selected_keypoints.append(keypoints[result_list[i]])

    # return selected_keypoints


class ProcessTOPM():
    def __init__(self,context) -> None:
        args = context.args
        self.num_ret_points = 50
        self.tolerance = 0.1
        self.image_size = args.resize
        self.args = args
        pass
    
    def run(self,feature_set):
        logging.debug("START TOPM")
        f_set2 = LocalFeatureSet(self.args)
        dataset_name = self.args.dataset_name
        f_set2._feature_dir = join(self.args.datasets_folder, dataset_name, "images","test","features_TOPM_" + timestr())
        os.mkdir(f_set2._feature_dir)
        result={}
        
        featureset = feature_set.get_features()
        for imgid,features in featureset.items():
            keypoints = features["keypoints"]
            #keypoints必须按照响应值倒序排序
            #暂时没有进行排序
            # x_index = np.lexsort(keypoints[:,2])
            # x_index = x_index[::-1]
            # keypoints = keypoints[x_index]
            
            selected_keypoints = keypoints[:self.num_ret_points ]
            selected_descriptors = features["descriptors"][:self.num_ret_points ]
            
            featuresnew = {}
            featuresnew["keypoints"] = selected_keypoints
            featuresnew["descriptors"] = selected_descriptors
            
            f_set2.add_features(imgid,featuresnew)
        f_set2.write()
        
        return f_set2

    def run_img(self,features):

            keypoints = features["keypoints"]
            # result_list = ssc(keypoints,self.num_ret_points,self.tolerance,self.image_size[1],self.image_size[0])
            
            # selected_keypoints = keypoints[result_list,:]
            # selected_descriptors = features["descriptors"][result_list,:]
            selected_keypoints = keypoints[:self.num_ret_points ]
            selected_descriptors = features["descriptors"][:self.num_ret_points ]


            featuresnew = {}
            featuresnew["keypoints"] = selected_keypoints
            featuresnew["descriptors"] = selected_descriptors

            return featuresnew